/*
 * author      : 吴中勤 
 * date        : 2019/3/11
 * description : TODO
 */

package org.wzq.learn.infrastructure.lib.trace;

import brave.Span;
import brave.Tracer;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.sleuth.instrument.web.TraceWebServletAutoConfiguration;
import org.springframework.core.annotation.Order;
import org.springframework.web.filter.GenericFilterBean;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

@Order(TraceWebServletAutoConfiguration.TRACING_FILTER_ORDER + 1)
public class FilterInitContext extends GenericFilterBean {

    private static final Logger logger = LoggerFactory.getLogger(FilterInitContext.class);

    private final Tracer tracer;

    public FilterInitContext(Tracer tracer) {
        this.tracer = tracer;
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        if (logger.isInfoEnabled()) logger.info("MyZipkinFilter filter");

        String xReqId = ((HttpServletRequest) request).getHeader("x-request-id");
        if (StringUtils.isNotBlank(xReqId)) {
            if (logger.isInfoEnabled()) logger.info("got x-request-id from request header");
            CustomThreadContext.setXRequestId(xReqId);
        }
        CustomThreadContext.setMock(checkIsMock((HttpServletRequest) request));

        Span currentSpan = this.tracer.currentSpan();
        if (currentSpan == null) {
            chain.doFilter(request, response);
            return;
        }

        // for readability we're returning trace id in a hex form
        String xTraceId = currentSpan.context().traceIdString();
        ((HttpServletResponse) response).addHeader("x-trace-id", xTraceId);
        currentSpan.tag("x-trace-id", xTraceId);
        CustomThreadContext.setXTraceId(xTraceId);
        // x-trace-id 不需要放到 appRuntimeEnv中

        try {
            chain.doFilter(request, response);
        } finally {
            CustomThreadContext.clean();
        }
    }

    private static final String KEY_FOR_MOCK = "mock";
    private static final String SYMBOL_AND = "&";
    private static final String SYMBOL_EQUAL = "=";

    // 直接使用 String rpi = request.getParameter("rpi"); 会破坏 stream
    private boolean checkIsMock(HttpServletRequest request) {

        // 检查请求头
        String mockInHeader = request.getHeader(KEY_FOR_MOCK);
        if (StringUtils.isNotEmpty(mockInHeader)) {
            return true;
        }

        // 检查query string
        String queryString = request.getQueryString();

        if (StringUtils.isEmpty(queryString)) {
            return false;
        }

        String[] parameters = queryString.split(SYMBOL_AND);

        for (String parameter : parameters) {
            String[] keyValuePair = parameter.split(SYMBOL_EQUAL);
            String key = keyValuePair[0];
            if (KEY_FOR_MOCK.equals(key)) {
                return true;
            }
        }

        return false;
    }
}